#TODO 导入相关的Python库或Python模块
import tkinter as tk
from tkinter import filedialog
import torch
from torch import nn
from torchvision import transforms
from PIL import Image, ImageOps, ImageDraw
import torch.nn.functional as F

from PIL import ImageTk 
import io

# 定义一个简单的卷积神经网络（CNN）模型，用于图像分类任务
class CNNNet(nn.Module):
    def __init__(self):
        super(CNNNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(500, 50)
        self.fc2 = nn.Linear(50, 4)

    # 前向传播
    def forward(self, x):
        x = torch.relu(torch.max_pool2d(self.conv1(x), 2))
        x = torch.relu(torch.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 500)
        x = torch.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

# 加载模型
model = CNNNet()
#TODO 本题文件的加载路径不唯一，需根据实际情况进行修改
model.load_state_dict(torch.load("D:\\横向1\\hengxiang\\2\\pre.pt"))
model.eval()

# 图像转换流程
transform = transforms.Compose([
    transforms.Resize((32, 32)),# 将图像大小调整为32像素宽和32像素高
    transforms.Grayscale(),     # 将彩色图像转换为灰度图像
    transforms.ToTensor(),      # 将PIL图像或numpy数组转换为torch.FloatTensor
    transforms.Normalize((0.1307,), (0.3081,)) #对图像像素值进行标准化的转换
])

# 预测图像里的字母
def predict_letter(image):
    image = ImageOps.invert(image.convert('L')) # 将图像转换为灰度模式并反转图像的颜色
    image = transform(image).unsqueeze(0)
    with torch.no_grad():
        output = model(image)
        _, predicted = torch.max(output, 1)
        letter = ['A', 'B', 'C', 'D'][predicted.item()]
    return letter

#TODO 创建一个简单的手写识别应用
class App:
    # 定义可视化界面功能
    # TODO 保存图片
    def save_image(self):
        # TODO（1） 调用filedialog模块中的asksaveasfilename函数实现弹出对话框，选择保存文件路径。
        file_path = filedialog.asksaveasfilename(defaultextension=".png")
        if file_path:
            # TODO（2） 保存图片到文件
            self.image.save(file_path)

    #TODO 字母识别
    def predict(self):
        #TODO（1） 调用predict_letter函数
        digit = predict_letter(self.image)
        self.label.config(text=f"识别结果：{digit}")

    # TODO 上传图片识别
    def upload(self):
        # TODO（1） 调用filedialog模块的askopenfilename函数实现弹出对话框，选择上传文件路径。
        file_path = filedialog.askopenfilename()
        if file_path:
            digit = predict_letter(Image.open(file_path))
            self.label.config(text=f"识别结果：{digit}")

    #TODO 在画布上绘制字母
    def paint(self, event):
        x1, y1 = (event.x - 1), (event.y - 1) #计算坐标，通过 event.x 和 event.y 获取鼠标的当前位置
        x2, y2 = (event.x + 1), (event.y + 1)
        # TODO（1） 调用Canvas组件的create_oval方法，绘制一个非常小的椭圆当作点，填充颜色为黑色，宽度为5
        self.canvas.create_oval(x1, y1, x2, y2, fill="black", width=5)
        #TODO（2） 绘制从 (x1, y1) 到 (x2, y2) 的线条，填充颜色为黑色，宽度为5
        self.canvas.create_line(x1, y1, x2, y2, fill="black", width=5)

    #TODO 清空画布和图像
    def clear(self):
        #TODO（1） 清空画布
        self.canvas.delete("all")
        self.image = Image.new("RGB", (300, 200), "white")
        self.draw = ImageDraw.Draw(self.image)

    # 初始化可视化界面
    def __init__(self, root):
        self.root = root
        self.root.title("手写字母识别")

        #TODO 创建图形用户界面（GUI）中的画布组件，在界面上进行绘图
        #TODO（1） 创建一个可视化界面窗口Canvas对象，宽300像素、高200像素、背景色为白色
        self.canvas = tk.Canvas(root, width=300, height=200, bg='white')
        self.canvas.pack()
        self.canvas.bind("<B1-Motion>", self.paint)

        #TODO（2） 设置‘识别’按钮
        self.predict_button = tk.Button(root, text="识别", command=self.predict)
        self.predict_button.pack()

        #TODO（3） 设置‘清空’按钮
        self.clear_button = tk.Button(root, text="清空", command=self.clear)
        self.clear_button.pack()

        #TODO（4） 设置‘保存图片’按钮
        self.save_button = tk.Button(root, text="保存图片", command=self.save_image)
        self.save_button.pack()

        #TODO（5） 设置‘导入图片’按钮
        self.upload_button = tk.Button(root, text="导入图片", command=self.upload)
        self.upload_button.pack()

        self.label = tk.Label(root, text="识别结果：")
        self.label.pack()

        #TODO（6） 调用Pillow库的Image.new方法初始化一个画布，大小为300x200像素，背景色为白色
        self.image = Image.new("RGB", (300, 200), "white")
        self.draw = ImageDraw.Draw(self.image)


if __name__ == "__main__":
    root = tk.Tk()
    app = App(root)
    root.mainloop()
